import numpy as np
import pandas as pd
from scipy.optimize import curve_fit
from pyswarm import pso
import pyswarms as ps
from pyswarms.utils.functions import single_obj as fx
from EUR_predict import build_ensemble_model, process_data, cross_validate_and_predict, eur_function
from sqlalchemy import create_engine, MetaData, Table
from sqlalchemy import text
from flask import Flask, request, jsonify
import uuid
from EUR_predict import build_ensemble_model, cross_validate_and_predict, eur_function
from sqlalchemy import create_engine
from sklearn.ensemble import RandomForestRegressor
from xgboost import XGBRegressor
from lightgbm import LGBMRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.linear_model import Lasso, Ridge
import math
from datetime import datetime
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import numpy as np
from sklearn.model_selection import GridSearchCV
from sklearn.svm import SVR
from math import sqrt
from sklearn.linear_model import HuberRegressor
from sklearn.preprocessing import MinMaxScaler
from sklearn.preprocessing import StandardScaler
import xgboost as xgb
from sklearn.model_selection import KFold
from sklearn.model_selection import RandomizedSearchCV
from sklearn.ensemble import StackingRegressor
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.impute import KNNImputer

username = 'root'
password = '123456'
host = '192.168.0.189'
port = '3306'
database = 'shale-gas'
'''
username = 'root'
password = '123456'
host = '10.63.211.120'
port = '3306'
database = 'mysql'
'''
engine = create_engine(f'mysql+pymysql://{username}:{password}@{host}:{port}/{database}')

# Establish DB connection
def get_db_connection():
    engine = create_engine(f'mysql+pymysql://{username}:{password}@{host}:{port}/{database}')
    return engine

from sklearn.preprocessing import LabelEncoder

# 定义类别变量编码函数
def encode_categorical_variables(X, label_encoder_type=None, label_encoder_region=None):
    """
    对输入数据框的类别变量进行编码。
    如果提供了现有的 LabelEncoder 实例，则使用这些实例进行编码。
    """
    if label_encoder_type is None:
        label_encoder_type = LabelEncoder()
        X['well_type'] = label_encoder_type.fit_transform(X['well_type'])
    else:
        X['well_type'] = label_encoder_type.transform(X['well_type'])

    if label_encoder_region is None:
        label_encoder_region = LabelEncoder()
        X['core_area'] = label_encoder_region.fit_transform(X['core_area'])
    else:
        X['core_area'] = label_encoder_region.transform(X['core_area'])

    return X, label_encoder_type, label_encoder_region

def process_data(df_combine):
    # 移除气井类型为空的数据
    df_combine = df_combine.dropna()
    gas_label = 'well_type'
    gas_label = df_combine.well_type
    df_combine = df_combine[pd.notna(gas_label)]
    df_combine.set_index('well_no', inplace=True)

    target_m = 'm_values'
    target_a = 'a_values'
    target_p = 'days330_first_year'

    del_vars = ['m_values', 'a_values', 'days330_first_year']
    X_all = df_combine.drop(del_vars, axis=1)

    # 提取目标变量
    y_all_a = df_combine[target_a]
    y_all_m = df_combine[target_m]
    y_all_p = df_combine[target_p]

    # 编码类别变量
    X_all, label_encoder_type, label_encoder_region = encode_categorical_variables(X_all)

    # 使用 KNN 填补缺失值
    imputer = KNNImputer(n_neighbors=3)
    imputed_data = imputer.fit_transform(X_all)

    # 转换为 DataFrame 并恢复索引
    X_imputed = pd.DataFrame(imputed_data, columns=X_all.columns)
    X_imputed['well_no'] = X_all.index
    X_imputed.set_index('well_no', inplace=True)

    # 处理后的目标变量
    y_process_a = df_combine[target_a]
    y_process_m = df_combine[target_m]
    y_process_p = df_combine[target_p]

    # 训练集的特征和标签
    Dataset_X = X_imputed[pd.notna(y_process_p.values)]
    Dataset_y_a = y_process_a[pd.notna(y_process_a.values)]
    Dataset_y_m = y_process_m[pd.notna(y_process_m.values)]
    Dataset_y_p = y_process_p[pd.notna(y_process_p.values)]

    return Dataset_X, Dataset_y_a, Dataset_y_m, Dataset_y_p, label_encoder_type, label_encoder_region


def get_processed_data():
    engine = get_db_connection()
    table_name = 'gas_well_para'
    noprocess_var = ['id', 'update_by', 'update_time', 'sys_org_code', 'create_by', 'create_time', 'actual_production',
                     'duong_production', 'lng', 'lat', 'well_state']
    df = pd.read_sql_table(table_name, con=engine)
    df = df.drop(noprocess_var, axis=1)
    df = df.dropna()

    predict_variables = ['thickness1_layer', 'brmc1_layer', 'toc1_layer',
                         'por1_layer', 'qall1_layer', 'sg1_layer', 'coefficient_pressure',
                         'subject_hop_crowd', 'subject_hop_hole', 'crowd_distance',
                         'af_stage_length', 'sand_strength', 'liquid_strength', 'dro4mbabo_rate',
                         'segment_length','drilling_depth']
    base_variables = ['well_no', 'well_type', 'm_values', 'a_values', 'days330_first_year', 'core_area']
    df = df[base_variables + predict_variables]

    # Data processing
    Dataset_X, Dataset_y_a, Dataset_y_m, Dataset_y_p, label_encoder_type, label_encoder_region = process_data(df)

    df_combined = pd.concat([Dataset_X, Dataset_y_a, Dataset_y_m, Dataset_y_p], axis=1)
    original_columns = list(Dataset_X.columns)
    new_columns = ['a_values', 'm_values', 'days330_first_year']
    df_combined.columns = original_columns + new_columns

    df_combined = df_combined[pd.notna(df_combined['a_values'].values)]
    df_combined = df_combined[pd.notna(df_combined['m_values'].values)]
    y_a = df_combined['a_values']
    y_m = df_combined['m_values']
    y_p = df_combined['days330_first_year']
    X_all = df_combined.drop(['a_values', 'm_values', 'days330_first_year'], axis=1)

    return X_all, y_a, y_m, y_p, label_encoder_type, label_encoder_region

# Define function to perform EUR calculation and optimization
def predict_para(newdata, scaler, model_a, model_m, model_p):
    new_scaled = scaler.transform(newdata)
    a_predict = model_a.predict(new_scaled)
    m_predict = model_m.predict(new_scaled)
    first_production = model_p.predict(new_scaled)
    return a_predict, m_predict, first_production

def eur_function(first_production, a_fit, m_fit, dur_eur=20):
    start_value = 331
    end_value = 330 * dur_eur
    dur = np.arange(start_value, end_value + 1)
    inital_q = intial_production_calculation(first_production, a_fit, m_fit, days=330)
    daily_pro = duong_decay((dur, np.repeat(inital_q, len(dur))), a_fit, m_fit)
    year_production = sum_year_production(daily_pro, 330)
    year_production_all = np.insert(year_production, 0, first_production)
    return year_production_all, sum(year_production_all)

def sum_year_production(data, n=330):
    sums = []
    i = 0
    while i < len(data):
        chunk = data[i:i + n]
        sums.append(sum(chunk))
        i += n
    return sums

def revenue_function(year_production, price=0.989, rate=0.06):
    sy = 0
    year_sy = []
    for j in range(len(year_production)):
        cl = year_production[j]
        sy = sy + ((cl * price) / ((1 + rate) ** j))
        year_sy.append((cl * price) / ((1 + rate) ** j))
        print(sy)
    return sy, year_sy

def cost_function(df, var, para, inter=2120.96):
    cost_df = df[var].values
    cost_all = np.dot(cost_df, para)[0] + inter
    return cost_all

#revise
def duong_decay(X, a, m):
    t = X[0]
    qi = X[1]
    return qi * t ** (-m) * np.exp( (a / (1 - m)) * (t ** (1 - m) - 1))

def intial_production_calculation(production, a, m, days):
    Q = production
    t = days
    return Q*a /(np.exp( (a / (1 - m)) * (t ** (1 - m) - 1)))

def cumulative_production(qi, a, m, t=330 * 6, t1=330):
    Q = (qi / a) * (np.exp((a / (1 - m)) * (t ** (1 - m) - 1)))
    Q1 = (qi / a) * (np.exp((a / (1 - m)) * (t1 ** (1 - m) - 1)))
    return Q - Q1
'''
def predict_para(newdata,scaler,model_a,model_m,model_p):
    new_scaled = scaler.transform(newdata)
    a_predict = model_a.predict(new_scaled)
    m_predict = model_m.predict(new_scaled)
    first_production = model_p.predict(new_scaled)

    return a_predict, m_predict, first_production'''

def eur_function(first_production, a_fit, m_fit, dur_eur=20):
    start_value = 331
    end_value = 330 * dur_eur
    dur = np.arange(start_value, end_value + 1)
    inital_q = intial_production_calculation(first_production, a_fit, m_fit, days=330)
    daily_pro = duong_decay((dur, np.repeat(inital_q, len(dur))), a_fit, m_fit)
    year_production = sum_year_production(daily_pro, 330)
    year_production_all = np.insert(year_production, 0, first_production)
    return year_production_all, sum(year_production_all)

def sum_year_production(data, n=330):
    sums = []
    i = 0
    while i < len(data):
        chunk = data[i:i + n]
        sums.append(sum(chunk))
        i += n
    return sums

def revenue_function(year_production, price=0.989, rate=0.06):
    sy = 0
    year_sy = []
    for j in range(len(year_production)):
        cl = year_production[j]
        sy = sy + ((cl * price) / ((1 + rate) ** j))
        year_sy.append((cl * price) / ((1 + rate) ** j))
    return sy, year_sy

def cost_function(df, var, para, inter=2120.96):
    cost_df = df[var].values
    cost_all = np.dot(cost_df, para)[0] + inter

    return cost_all

business_cost = [0, 231.24, 365.31, 307.36, 250.72, 271.33, 220.50, 190.42, 160.69, 147.23, 137.79, 130.93,
                 125.40, 120.33, 115.65
    , 111.35, 107.40, 103.76, 100.41, 97.33]

def eur_new_function(first_production, ratio):
    cl = first_production.copy()
    num_years = len(ratio)
    rc = []
    for i in range(num_years):
        rc=np.append(rc, cl, axis=0)
        cl *= (1 - ratio[i])
    rc=np.append(rc, cl, axis=0)
    return rc, sum(rc)

def revenue_new_function(rc, ratio, qj=0.989, business = business_cost):
    num_years = len(ratio)
    sy = 0
    year_sy = []
    for j in range((num_years)):
        cl = rc[j]
        sy = sy + ( ((cl*qj) - business[j])/((1+0.06)**(j)) )
        year_sy.append(( ((cl*qj) - business[j])/((1+0.06)**(j)) ))
    sy = sy + (((cl * qj) - business[j+1]) / ((1 + 0.06) ** (j+1)))
    year_sy.append(((cl * qj) - business[j+1]) / ((1 + 0.06) ** (j+1)))
    return sy,  year_sy

def cost_new_function(df, var, para, inter=2873.7):
    cost_new = df.copy()
    cost_new['liquid_strength'] = df['liquid_strength'] * df['af_stage_length']
    cost_new['sand_strength'] = df['sand_strength'] * df['af_stage_length']

    cost_df = cost_new[var].values
    cost_all = np.dot(cost_df, para)[0] + inter

    return cost_all

def main():
    # 单井建设成本(万元) = 2120.96 -23.66 * 簇间距(米) + 0.64 * 主体单段簇数 + 34.28 * 用液强度(m3/m) + 12.89 * 加砂强度(t/m) -21.62 * 分段段长(m) + 0.55 * 完井深度(m) + 1.41 * 压裂断长(m)
    pass

if __name__ == '__main__':
    main()

    #单井建设成本(万元) = 2120.96 -23.66 * 簇间距(米) + 0.64 * 主体单段簇数 + 34.28 * 用液强度(m3/m) + 12.89 * 加砂强度(t/m) -21.62 * 分段段长(m) + 0.55 * 完井深度(m) + 1.41 * 压裂断长(m)